from scapy.all import rdpcap, IP, TCP, Raw
from io import StringIO

def reassemble_tcp_streams(file_path):
    packets = rdpcap(file_path)
    flows = {}

    for pkt in packets:
        if IP in pkt and TCP in pkt and Raw in pkt:
            ip = pkt[IP]
            tcp = pkt[TCP]
            raw = pkt[Raw]

            # 五元组定义流
            flow_key = (
                ip.src,
                ip.dst,
                tcp.sport,
                tcp.dport,
                'TCP'
            )

            reverse_flow_key = (
                ip.dst,
                ip.src,
                tcp.dport,
                tcp.sport,
                'TCP'
            )

            # 判断方向
            payload = raw.load
            timestamp = pkt.time

            # 按照流方向组织数据
            if flow_key in flows:
                flows[flow_key].append(('->', payload, timestamp))
            elif reverse_flow_key in flows:
                flows[reverse_flow_key].append(('<-', payload, timestamp))
            else:
                flows[flow_key] = [('->', payload, timestamp)]

    buffer = StringIO()
    # 打印每个流的内容
    for idx, (flow_key, stream) in enumerate(flows.items()):
        #print(f"\n=== TCP Stream #{idx} ===")
        #print(f"Flow: {flow_key[0]}:{flow_key[2]} -> {flow_key[1]}:{flow_key[3]}")

        client_data = b''
        server_data = b''

        for direction, payload, ts in stream:
            if direction == '->':
                client_data += payload
            else:
                server_data += payload

        try:
            #print("\n--- Client to Server ---")
            #print(client_data.decode('utf-8', errors='replace'))
            buffer.write(client_data.decode('utf-8', errors='replace'))
            #print("\n--- Server to Client ---")
            #print(server_data.decode('utf-8', errors='replace'))
            buffer.write(server_data.decode('utf-8', errors='replace'))
        except Exception as e:
            print(f"[无法解码载荷: {e}]")

        #print("=" * 50)

    #return flows
    print(buffer.getvalue())
    return buffer.getvalue()

# 调用函数分析 pcap 文件
reassemble_tcp_streams("E:\\13.20230718.00080.pcap")